import ctypes
import os
import platform
import numpy as np
from typing import List, Tuple

class QuantumGateOp(ctypes.Structure):
    _fields_ = [
        ("type", ctypes.c_int),
        ("target", ctypes.c_int),
        ("control", ctypes.c_int),
        ("parameter", ctypes.c_double)
    ]

class QuantumSimulator:
    def __init__(self):
        # 加载动态库
        system = platform.system()
        if system == "Windows":
            self.lib = ctypes.CDLL("quantum.dll")
        elif system == "Linux":
            self.lib = ctypes.CDLL("/home/xxieww_group/ustcjyd/huaweiLinkeq/backend/libquantum.so")
        else:
            raise OSError("Unsupported operating system")
        
        # 设置函数原型
        self.lib.simulate_quantum_circuit.argtypes = [
            ctypes.c_int,                    # num_qubits
            ctypes.c_int,                    # num_operations
            ctypes.POINTER(QuantumGateOp),   # operations
            ctypes.POINTER(ctypes.c_int)     # state_size (output)
        ]
        self.lib.simulate_quantum_circuit.restype = ctypes.POINTER(ctypes.c_double)
        
        self.lib.free_quantum_state.argtypes = [ctypes.POINTER(ctypes.c_double)]
        self.lib.free_quantum_state.restype = None
    
    def simulate(
        self,
        num_qubits: int,
        operations: List[Tuple[str, List[int], float]]
    ) -> np.ndarray:
        """
        执行量子电路模拟
        
        参数:
            num_qubits: 量子比特数
            operations: 量子门操作列表
                格式: [(门类型, [作用比特], 参数)]
                示例: [("H", [0], 0), ("CNOT", [0, 1], 0)]
        
        返回:
            状态向量 (复数numpy数组)
        """
        # 门类型映射
        gate_map = {
            "H": 0,
            "X": 1,
            "Y": 2,
            "Z": 3,
            "CNOT": 4,
            "SWAP": 5,
            "RX": 6,
            "RY": 7,
            "RZ": 8
        }
        
        # 准备操作数组
        c_ops = (QuantumGateOp * len(operations))()
        for i, op in enumerate(operations):
            gate_type, qubits, param = op
            c_ops[i].type = gate_map[gate_type]
            c_ops[i].parameter = param
            
            if len(qubits) >= 1:
                c_ops[i].target = qubits[0]
            if len(qubits) >= 2:
                c_ops[i].control = qubits[1]
            else:
                c_ops[i].control = -1  # 表示未使用
        
        # 调用C++函数
        state_size = ctypes.c_int(0)
        result_ptr = self.lib.simulate_quantum_circuit(
            num_qubits,
            len(operations),
            c_ops,
            ctypes.byref(state_size)
        )
        
        # 处理结果
        if not result_ptr:
            raise RuntimeError("Quantum simulation failed")
        
        try:
            # 转换为numpy数组
            float_array = np.ctypeslib.as_array(
                result_ptr, 
                shape=(state_size.value,)
            )
            
            # 转换为复数数组
            real = float_array[0::2]
            imag = float_array[1::2]
            state_vector = real + 1j * imag
            
            return state_vector
        finally:
            # 确保释放内存
            self.lib.free_quantum_state(result_ptr)

if __name__ == "__main__":
    # 使用示例
    simulator = QuantumSimulator()
    
    # 创建Bell态电路: H(0) -> CNOT(0,1)
    num_qubits = 3
    operations = [
        ("H", [0], 0.0),     # 在0号比特应用Hadamard门
        ("CNOT", [0, 1], 0.0), # CNOT门，控制0，目标1
        ("RZ", [2], 1.0)   # 在0号比特应用Hadamard门
    ]
    print(type(operations))
    
    # 运行模拟
    state_vector = simulator.simulate(num_qubits, operations)
    
    # 打印结果
    print("量子态向量:")
    for i, amp in enumerate(state_vector):
        print(f"|{i:0{num_qubits}b}>: {amp.real:.4f}{amp.imag:+.4f}i")
    
    # 验证Bell态
    print("\n验证Bell态 (|00> + |11>)/√2:")
    print(f"|00>振幅: {state_vector[0]}")
    print(f"|11>振幅: {state_vector[3]}")